from typing import List
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.schema import BaseOutputParser
import re
from PyCmpltrtok.common import sep
from langchain.output_parsers import CommaSeparatedListOutputParser

llm = OpenAI(
    streaming=True,
    verbose=True,
    callbacks=[],
    openai_api_key="token1",
    openai_api_base=f"http://127.0.0.1:6001/v1",
    model_name="chatglm2-6b-int4",
    temperature=0.5,
    openai_proxy=None,
    top_p=1.0,
)

template = """你是一个生成逗号分隔列表的助手。
你生成"{category}"这个类别的5个对象组成的逗号分隔列表。
只生成逗号分隔列表，不生成其他内容。
"""

prompt_tpl = PromptTemplate.from_template(template)

category = "汽车"  # colors, presidents, 颜色, 汽车, 各种汽车

sep('prompt | llm | parser')
chain = prompt_tpl | llm | CommaSeparatedListOutputParser()
result = chain.stream({"category": category})
print(result)
for i, x in enumerate(result):
    print(f'{i}: |{x}|')
# >> ['red', 'blue', 'green', 'yellow', 'orange']

sep('prompt | llm')
chain = prompt_tpl | llm
result = chain.stream({"category": category})
print(result)
for i, x in enumerate(result):
    print(f'{i}: |{x}|')
